use num_bigint::BigUint;
use num_rational::Ratio;
use std::cmp::Ordering;
use std::collections::HashMap;

/// Nonterminal symbols
pub type N = u32;

pub type RHS = Vec<Symbol>;

/// Symbols occurring in a Sentential Form
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Symbol {
    Nonterminal(N),
    Terminals(Vec<u8>),
}

/// A number $\frac{k}{2^f}$ where $k \in \mathbb{N}$ and $f \in \mathbb{Z}$
#[derive(Clone, Debug)]
pub struct BinaryFraction {
    /// Numerator (>0)
    pub k: BigUint,
    /// Denominator power.
    /// f<0 means k is shifted to the left
    /// f>0 means k is shifted to the right
    // Is the whole range of u64 realistically needed? Would anything be gained by restricting?
    pub f: i64,
}

pub type Prob = BinaryFraction;
pub fn mk_prob(numer: BigUint, logdenom: u64) -> Prob {
    BinaryFraction::new(numer, logdenom as i64)
}

use core::fmt;
impl fmt::Display for BinaryFraction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.f == 0 {
            write!(f, "{}", self.k)
        } else {
            write!(f, "{} / 2^-{}", self.k, self.f)
        }
    }
}

use core::ops::{Add, Mul, Sub};
use num_traits::Zero;

impl BinaryFraction {
    #[inline]
    pub fn new(k: BigUint, f: i64) -> BinaryFraction {
        let mut ret = BinaryFraction::new_raw(k, f);
        ret.reduce();
        ret
    }

    /// @precondition: r.denom() must be a power of 2
    pub fn from_bigrational(r: Ratio<BigUint>) -> BinaryFraction {
        assert!(r.denom() >> r.denom().trailing_zeros().unwrap() == One::one());
        BinaryFraction::new(
            r.numer().clone(),
            r.denom().trailing_zeros().unwrap() as i64,
        )
    }

    pub fn from_float_between_0_and_1(f: f32) -> Option<BinaryFraction> {
        let b = Ratio::from_float(f)?;
        if !(b >= Zero::zero() && b <= One::one()) {
            return None;
        }
        Some(BinaryFraction::from_bigrational(Ratio::new_raw(
            b.numer().to_biguint()?,
            b.denom().to_biguint()?,
        )))
    }

    pub fn into_bigrational(r: Self) -> Ratio<BigUint> {
        let one: BigUint = One::one();
        if r.f < 0 {
            Ratio::new(r.k << -r.f, one)
        } else {
            Ratio::new(r.k, one << r.f)
        }
    }

    pub fn new_raw(k: BigUint, f: i64) -> BinaryFraction {
        BinaryFraction { k, f }
    }

    #[inline]
    pub fn from_biguint(k: BigUint) -> BinaryFraction {
        BinaryFraction::new(k, 0)
    }

    /// Puts self into lowest terms, with k odd,
    /// 2^{-f} = lowest non-zero bit.
    /// If k = 0, f = 0.
    fn reduce(&mut self) {
        if self.k.is_zero() {
            self.f = 0;
            return;
        }
        match self.k.trailing_zeros() {
            None => {}
            Some(z) => {
                self.f -= z as i64;
                self.k >>= z;
            }
        }
    }

    /// Return a reduced copy of itself
    pub fn reduced(&self) -> BinaryFraction {
        let mut ret = self.clone();
        ret.reduce();
        ret
    }

    /// Rounds towards zero.
    pub fn trunc_as_biguint(&self) -> BigUint {
        if self.f > 0 {
            self.k.clone() >> self.f
        } else {
            self.k.clone() << -self.f
        }
    }

    #[inline]
    pub fn trunc(&self) -> BinaryFraction {
        BinaryFraction::from_biguint(self.trunc_as_biguint())
    }
}

mod opassign {
    use core::ops::{AddAssign, MulAssign, SubAssign};
    use std::cmp::Ordering;

    use crate::BinaryFraction;

    impl AddAssign for BinaryFraction {
        fn add_assign(&mut self, other: BinaryFraction) {
            let d = self.f - other.f;

            match d.cmp(&0) {
                Ordering::Equal => self.k += other.k,
                // self.f < other.f i.e. other is further to the right
                Ordering::Less => {
                    self.f = other.f;
                    self.k <<= -d;
                    self.k += other.k;
                }
                // self.f > other.f i.e. other is further to the left
                Ordering::Greater => {
                    self.k += other.k << d;
                }
            }

            self.reduce();
        }
    }

    impl SubAssign for BinaryFraction {
        fn sub_assign(&mut self, other: BinaryFraction) {
            let d = self.f - other.f;

            match d.cmp(&0) {
                Ordering::Equal => self.k -= other.k,
                // self.f < other.f i.e. other is further to the right
                Ordering::Less => {
                    self.f = other.f;
                    self.k <<= -d;
                    self.k -= other.k;
                }
                // self.f > other.f i.e. other is further to the left
                Ordering::Greater => {
                    self.k -= other.k << d;
                }
            }

            self.reduce();
        }
    }

    impl MulAssign for BinaryFraction {
        fn mul_assign(&mut self, other: BinaryFraction) {
            self.f += other.f;
            self.k *= other.k;
            self.reduce();
        }
    }
}

impl Add<BinaryFraction> for BinaryFraction {
    type Output = BinaryFraction;

    #[inline]
    fn add(self, other: BinaryFraction) -> BinaryFraction {
        let mut result = self;
        result += other;
        result
    }
}

impl<'a> Add<&'a BinaryFraction> for BinaryFraction {
    type Output = BinaryFraction;

    #[inline]
    fn add(self, other: &BinaryFraction) -> BinaryFraction {
        self.add(other.clone())
    }
}

impl Sub<BinaryFraction> for BinaryFraction {
    type Output = BinaryFraction;

    #[inline]
    fn sub(self, other: BinaryFraction) -> BinaryFraction {
        let mut result = self;
        result -= other;
        result
    }
}

impl<'a> Sub<&'a BinaryFraction> for BinaryFraction {
    type Output = BinaryFraction;

    #[inline]
    fn sub(self, other: &BinaryFraction) -> BinaryFraction {
        self.sub(other.clone())
    }
}

impl Mul<BinaryFraction> for BinaryFraction {
    type Output = BinaryFraction;

    #[inline]
    fn mul(self, other: BinaryFraction) -> BinaryFraction {
        let mut result = self;
        result *= other;
        result
    }
}

impl<'a> Mul<&'a BinaryFraction> for BinaryFraction {
    type Output = BinaryFraction;

    #[inline]
    fn mul(self, other: &BinaryFraction) -> BinaryFraction {
        self.mul(other.clone())
    }
}

impl Mul<usize> for BinaryFraction {
    type Output = BinaryFraction;

    #[inline]
    fn mul(self, other: usize) -> BinaryFraction {
        self.mul(BinaryFraction::from_biguint(BigUint::from(other)))
    }
}

impl Zero for BinaryFraction {
    #[inline]
    fn zero() -> BinaryFraction {
        BinaryFraction {
            k: Zero::zero(),
            f: 0,
        }
    }

    #[inline]
    fn set_zero(&mut self) {
        self.k.set_zero();
        self.f = 0;
    }

    #[inline]
    fn is_zero(&self) -> bool {
        self.k.is_zero()
    }
}

impl One for BinaryFraction {
    #[inline]
    fn one() -> BinaryFraction {
        BinaryFraction {
            k: One::one(),
            f: 0,
        }
    }

    #[inline]
    fn set_one(&mut self) {
        self.k.set_one();
        self.f = 0;
    }

    #[inline]
    fn is_one(&self) -> bool {
        if self.f < 0 {
            (self.k.clone() >> -self.f).is_one()
        } else {
            (self.k.clone() << self.f).is_one()
        }
    }
}

use core::cmp;
impl Ord for BinaryFraction {
    #[inline]
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        // Is there a quicker way to perform the comparison?
        let d = self.f - other.f;

        match d.cmp(&0) {
            Ordering::Equal => self.k.cmp(&other.k),
            Ordering::Less => (self.k.clone() << -d).cmp(&other.k),
            Ordering::Greater => self.k.cmp(&(&other.k << d)),
        }
    }
}

impl PartialOrd for BinaryFraction {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for BinaryFraction {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == cmp::Ordering::Equal
    }
}

impl Eq for BinaryFraction {}

#[derive(Clone, Debug)]
pub struct RuleEntry {
    pub probability: Prob,
    pub cumul_lower: Prob,
    pub rhs: RHS,
}

/// Rule RHSes are *ordered* lists decorated with cumulative probabilities
/// Start symbol is 0
#[derive(Debug)]
pub struct SCFG {
    pub rules: HashMap<N, Vec<RuleEntry>>,
}

pub static START_SYMBOL: N = 0;

use num_traits::float::FloatCore;
use num_traits::identities::One;
use num_traits::FromPrimitive;
// Adapted from https://docs.rs/num-rational/0.3.0/src/num_rational/lib.rs.html#263-282

/// Converts a float into a rational number.
pub fn from_float<T: FloatCore>(f: T) -> Option<Ratio<BigUint>> {
    if !f.is_finite() {
        return None;
    }
    let (mantissa, exponent, _sign) = f.integer_decode();
    if exponent < 0 {
        let one: BigUint = One::one();
        let denom: BigUint = one << ((-exponent) as usize);
        let numer: BigUint = FromPrimitive::from_u64(mantissa).unwrap();
        Some(Ratio::new(numer, denom))
    } else {
        let mut numer: BigUint = FromPrimitive::from_u64(mantissa).unwrap();
        numer <<= exponent as usize;
        Some(Ratio::from_integer(numer))
    }
}

/// Probability lower and upper bounds,
/// with event probability expressed relatively to l(ower_bound) instead of absolutely,
/// as this will often save a bignum operation.
#[derive(Debug, Clone)]
pub struct BoundsLP {
    pub l: Prob,
    pub p: Prob,
}
